import io
import wave
import numpy as np
from pydub import AudioSegment
import simpleaudio as sa


def play_mp3(file_path):
    """播放MP3文件"""
    try:
        # 加载音频文件
        audio = AudioSegment.from_mp3(file_path)
        # 将 AudioSegment 转换为 raw audio data
        raw_data = audio.raw_data

        # 创建音频播放实例
        play_obj = sa.play_buffer(
            raw_data,
            num_channels=audio.channels,
            bytes_per_sample=audio.sample_width,
            sample_rate=audio.frame_rate
        )

        # 等待音频播放完毕
        play_obj.wait_done()
    except Exception as e:
        print("An error occurred:", e)


def play_wav(file_path):
    """播放WAV文件"""
    try:
        # 加载 WAV 文件
        wave_obj = sa.WaveObject.from_wave_file(file_path)

        # 播放音频
        play_obj = wave_obj.play()

        # 等待播放结束
        play_obj.wait_done()
    except Exception as e:
        print("An error occurred:", e)


def play_pcm(file_path, sample_rate=44100, channels=1, bit_depth=16, endian='big'):
    """
    播放一个 PCM 文件。

    参数:
    file_path (str): PCM 文件的路径
    sample_rate (int): 采样率，默认为 44100 Hz
    channels (int): 声道数，默认为单声道（1）
    bit_depth (int): 位深度，默认为 16 位
    endian (str): 字节序，默认为 'big'（大端模式），可选 'little'（小端模式）
    """
    try:
        audio = readPCM(file_path, sample_rate=sample_rate, channels=channels, sample_width=int(bit_depth/8))
        playMP3(audio)
    except Exception as e:
        pass


def play_bin(file_path, sample_rate=22050, channels=1, bit_depth=16, endian='big'):
    """
    播放一个 PCM 文件。

    参数:
    file_path (str): PCM 文件的路径
    sample_rate (int): 采样率，默认为 44100 Hz
    channels (int): 声道数，默认为单声道（1）
    bit_depth (int): 位深度，默认为 16 位
    endian (str): 字节序，默认为 'big'（大端模式），可选 'little'（小端模式）
    """
    try:
        audio = readBIN(file_path, sample_rate=sample_rate, channels=channels, sample_width=int(bit_depth/8))
        playMP3(audio)
    except Exception as e:
        pass

def convert_to_mp3(input_file, output_file, channels=1):
    """将任意音频文件转化成MP3格式"""
    try:
        # 加载音频文件
        audio = AudioSegment.from_file(input_file)

        # 转换为单声道
        audio_mono = audio.set_channels(channels)

        # 导出为 MP3 文件
        audio_mono.export(output_file, format="mp3")

        print(f"转换成功！输出文件为: {output_file}")
        return 0
    except Exception as e:
        print(f"转换失败：{e}")
        return -1


def convert_mp3_to_pcm(mp3_file_path: str, pcm_file_path: str, gain_db=0, hz=0):
    """
    MP3转PCM格式
    :param mp3_file_path:
    :param pcm_file_path:
    :param gain_db:
    :param hz: 采样率
    :return:
    """
    try:
        # 加载 MP3 文件
        audio = AudioSegment.from_mp3(mp3_file_path)

        # 配置增益
        audio = audio.apply_gain(gain_db)

        # 转换为单声道
        audio = audio.set_channels(1)

        if hz == 0:     # 设置采样率为 44100 Hz
            audio = audio.set_frame_rate(44100)
        elif hz == 1:   # 设置采样率为 22050 Hz
            audio = audio.set_frame_rate(22050)
        else:           # 默认设置采样率为 44100 Hz
            audio = audio.set_frame_rate(44100)

        # 转换为 16 位 PCM
        audio = audio.set_sample_width(2)  # 16 bits = 2 bytes

        # 获取音频数据
        pcm_data = np.array(audio.get_array_of_samples())

        # 将数据转换为 big-endian 格式
        pcm_data_big_endian = pcm_data.astype(np.int16).byteswap()

        # 写入 PCM 文件
        with open(pcm_file_path, 'wb') as pcm_file:
            pcm_file.write(pcm_data_big_endian)

        return 0
    except Exception as e:
        print(f"错误: {e}")
        return -1


def convert_mp3_to_bin(mp3_file_path: str, pcm_file_path: str, gain_db=0, hz=1):
    """
    MP3转PCM格式
    :param mp3_file_path:
    :param pcm_file_path:
    :param gain_db:
    :param hz: 采样率
    :return:
    """
    try:
        # 加载 MP3 文件
        audio = AudioSegment.from_mp3(mp3_file_path)

        # 配置增益
        audio = audio.apply_gain(gain_db)

        # 转换为单声道
        audio = audio.set_channels(1)

        if hz == 0:     # 设置采样率为 44100 Hz
            audio = audio.set_frame_rate(44100)
        elif hz == 1:   # 设置采样率为 22050 Hz
            audio = audio.set_frame_rate(22050)
        else:           # 默认设置采样率为 44100 Hz
            audio = audio.set_frame_rate(44100)

        # 转换为 16 位 PCM
        audio = audio.set_sample_width(2)  # 16 bits = 2 bytes

        # 获取音频数据
        pcm_data = np.array(audio.get_array_of_samples())

        # 将数据转换为 big-endian 格式
        pcm_data_big_endian = pcm_data.astype(np.int16).byteswap()

        # 写入 PCM 文件
        with open(pcm_file_path, 'wb') as pcm_file:
            # 写入 160 个 0x00 字节
            pcm_file.write(b'\x00' * 160)
            pcm_file.write(pcm_data_big_endian)

        return 0
    except Exception as e:
        print(f"错误: {e}")
        return -1



####### MP3 处理
def get_MP3_length(file_path):
    """读取MP3/WAV文件的时间长度（单位毫秒）"""
    # 加载音频文件
    audio = AudioSegment.from_file(file_path)
    # 计算音频长度（毫秒）
    duration = len(audio)
    return duration


def readMP3(file_path):
    """读取MP3/WAV文件的时间长度（单位毫秒）"""
    # 加载音频文件
    audio = AudioSegment.from_file(file_path)
    return audio


def readSpeedMP3(audio, factor):
    """对MP3倍速"""
    new_audio = audio._spawn(audio.raw_data, overrides={
        "frame_rate": int(audio.frame_rate * factor)
    })
    return new_audio.set_frame_rate(audio.frame_rate)


def readGainMP3(audio, gain):
    """对MP3音量增益"""
    # 注意：pydub的+dB和-dB是以10为底的对数刻度，所以+3dB实际上是音量的两倍，而-3dB是音量的一半。
    # 但是，pydub的volume方法接受的是一个线性增益因子，而不是分贝值。
    # 为了达到分贝增益的效果，我们需要使用10**(dB/20)来计算增益因子。
    return audio.apply_gain(gain)


def getMP3data(audio):
    # 创建一个BytesIO对象来模拟文件
    wav_buffer = io.BytesIO()
    # 将AudioSegment对象导出到BytesIO对象中，格式为WAV
    audio.export(wav_buffer, format="wav")
    # 将文件指针重置到BytesIO对象的开始位置，以便后续读取
    wav_buffer.seek(0)

    wave_file = wave.Wave_read(wav_buffer)

    num_channels = wave_file.getnchannels()
    sample_width = wave_file.getsampwidth()
    framerate = wave_file.getframerate()
    num_frames = wave_file.getnframes()

    # 读取音频帧数据
    frames = wave_file.readframes(num_frames)

    audio_data = np.frombuffer(frames, dtype=np.int16)

    # 如果是立体声，我们需要处理两个声道的数据
    if num_channels == 2:
        audio_data = np.reshape(audio_data, (-1, 2))
        # 这里我们只绘制一个声道（例如左声道）的波形图
        audio_data = audio_data[:, 0]

    return audio_data / 32768.0     # [-1, 1] 归一化处理


def getMP3_frame_rate(audio):
    return audio.frame_rate


def playMP3(audio):
    try:
        # 将 AudioSegment 转换为 raw audio data
        raw_data = audio.raw_data

        # 创建音频播放实例
        play_obj = sa.play_buffer(
            raw_data,
            num_channels=audio.channels,
            bytes_per_sample=audio.sample_width,
            sample_rate=audio.frame_rate
        )

        # 等待音频播放完毕
        play_obj.wait_done()
    except Exception as e:
        print("An error occurred:", e)


def getMP3Size(audio):
    # 创建一个 BytesIO 对象，写入一些数据
    mp3_buffer = io.BytesIO()
    # 将AudioSegment对象导出到BytesIO对象中，格式为WAV
    audio.export(mp3_buffer, format="mp3")
    size_in_bytes = mp3_buffer.getbuffer().nbytes
    return size_in_bytes


def outputMP3(audio, outputFile):
    try:
        # 导出为 MP3 文件
        audio.export(outputFile, format="mp3")
        # print(f"转换成功！输出文件为: {outputFile}")
        return 0
    except Exception as e:
        print(f"转换失败：{e}")
        return -1


####### PCM 处理
def readPCM(file_path, sample_rate=44100, channels=1, sample_width=2):
    """
    读取大端模式的 PCM 文件，并转换为 pydub.AudioSegment。

    :param file_path: PCM 文件路径
    :param sample_rate: PCM 文件的采样率（Hz）
    :param channels: 声道数（1: 单声道, 2: 立体声）
    :param sample_width: 样本宽度（字节数，例如 2 字节 = 16 位）
    :return: 一个 pydub.AudioSegment 对象
    """
    with open(file_path, "rb") as pcm_file:
        big_endian_data = pcm_file.read()

    # 将大端序的 PCM 数据转换为小端序
    # 每个样本的字节数由 sample_width 决定
    little_endian_data = bytearray()
    for i in range(0, len(big_endian_data), sample_width):
        # 读取每个样本，并反转字节顺序
        sample = big_endian_data[i:i + sample_width]
        little_endian_data.extend(sample[::-1])

    # 使用 pydub 读取小端序的 PCM 数据
    audio = AudioSegment(
        data=bytes(little_endian_data),
        sample_width=sample_width,
        frame_rate=sample_rate,
        channels=channels
    )
    return audio


def readSpeedPCM(audio, factor):
    """对MP3倍速"""
    new_audio = audio._spawn(audio.raw_data, overrides={
        "frame_rate": int(audio.frame_rate * factor)
    })
    return new_audio.set_frame_rate(audio.frame_rate)


def getPCMdata(audio):
    # 创建一个BytesIO对象来模拟文件
    wav_buffer = io.BytesIO()
    # 将AudioSegment对象导出到BytesIO对象中，格式为WAV
    audio.export(wav_buffer, format="wav")
    # 将文件指针重置到BytesIO对象的开始位置，以便后续读取
    wav_buffer.seek(0)

    wave_file = wave.Wave_read(wav_buffer)

    num_channels = wave_file.getnchannels()
    sample_width = wave_file.getsampwidth()
    framerate = wave_file.getframerate()
    num_frames = wave_file.getnframes()

    # 读取音频帧数据
    frames = wave_file.readframes(num_frames)

    audio_data = np.frombuffer(frames, dtype=np.int16)

    # 如果是立体声，我们需要处理两个声道的数据
    if num_channels == 2:
        audio_data = np.reshape(audio_data, (-1, 2))
        # 这里我们只绘制一个声道（例如左声道）的波形图
        audio_data = audio_data[:, 0]

    return audio_data / 32768.0     # [-1, 1] 归一化处理


def readGainPCM(audio, gain):
    """对MP3音量增益"""
    # 注意：pydub的+dB和-dB是以10为底的对数刻度，所以+3dB实际上是音量的两倍，而-3dB是音量的一半。
    # 但是，pydub的volume方法接受的是一个线性增益因子，而不是分贝值。
    # 为了达到分贝增益的效果，我们需要使用10**(dB/20)来计算增益因子。
    return audio.apply_gain(gain)


def getPCMSize(audio):
    # 创建一个 BytesIO 对象，写入一些数据
    pcm_buffer = io.BytesIO()

    # 获取音频数据
    pcm_data = np.array(audio.get_array_of_samples())

    # 将数据转换为 big-endian 格式
    pcm_data_big_endian = pcm_data.astype(np.int16).byteswap()

    # pcm_data_big_endian数据写入 PCM 文件
    pcm_buffer.write(pcm_data_big_endian.tobytes())

    size_in_bytes = pcm_buffer.getbuffer().nbytes
    return size_in_bytes


def outputPCM(audio, outputFile):
    try:
        # 获取音频数据
        pcm_data = np.array(audio.get_array_of_samples())

        # 将数据转换为 big-endian 格式
        pcm_data_big_endian = pcm_data.astype(np.int16).byteswap()

        # 写入 PCM 文件
        with open(outputFile, 'wb') as pcm_file:
            pcm_file.write(pcm_data_big_endian)

        return 0
    except Exception as e:
        print(f"错误: {e}")
        return -1


####### BIN 处理
def readBIN(file_path, sample_rate=22050, channels=1, sample_width=2):
    """
    读取大端模式的 PCM 文件，并转换为 pydub.AudioSegment。

    :param file_path: PCM 文件路径
    :param sample_rate: PCM 文件的采样率（Hz）
    :param channels: 声道数（1: 单声道, 2: 立体声）
    :param sample_width: 样本宽度（字节数，例如 2 字节 = 16 位）
    :return: 一个 pydub.AudioSegment 对象
    """
    with open(file_path, "rb") as pcm_file:
        pcm_file.seek(160)
        big_endian_data = pcm_file.read()

    # 将大端序的 PCM 数据转换为小端序
    # 每个样本的字节数由 sample_width 决定
    little_endian_data = bytearray()
    for i in range(0, len(big_endian_data), sample_width):
        # 读取每个样本，并反转字节顺序
        sample = big_endian_data[i:i + sample_width]
        little_endian_data.extend(sample[::-1])

    # 使用 pydub 读取小端序的 PCM 数据
    audio = AudioSegment(
        data=bytes(little_endian_data),
        sample_width=sample_width,
        frame_rate=sample_rate,
        channels=channels
    )
    return audio


def readSpeedBIN(audio, factor):
    """对MP3倍速"""
    new_audio = audio._spawn(audio.raw_data, overrides={
        "frame_rate": int(audio.frame_rate * factor)
    })
    return new_audio.set_frame_rate(audio.frame_rate)


def getBINdata(audio):
    # 创建一个BytesIO对象来模拟文件
    wav_buffer = io.BytesIO()
    # 将AudioSegment对象导出到BytesIO对象中，格式为WAV
    audio.export(wav_buffer, format="wav")
    # 将文件指针重置到BytesIO对象的开始位置，以便后续读取
    wav_buffer.seek(0)

    wave_file = wave.Wave_read(wav_buffer)

    num_channels = wave_file.getnchannels()
    sample_width = wave_file.getsampwidth()
    framerate = wave_file.getframerate()
    num_frames = wave_file.getnframes()

    # 读取音频帧数据
    frames = wave_file.readframes(num_frames)

    audio_data = np.frombuffer(frames, dtype=np.int16)

    # 如果是立体声，我们需要处理两个声道的数据
    if num_channels == 2:
        audio_data = np.reshape(audio_data, (-1, 2))
        # 这里我们只绘制一个声道（例如左声道）的波形图
        audio_data = audio_data[:, 0]

    return audio_data / 32768.0     # [-1, 1] 归一化处理


def readGainBIN(audio, gain):
    """对MP3音量增益"""
    # 注意：pydub的+dB和-dB是以10为底的对数刻度，所以+3dB实际上是音量的两倍，而-3dB是音量的一半。
    # 但是，pydub的volume方法接受的是一个线性增益因子，而不是分贝值。
    # 为了达到分贝增益的效果，我们需要使用10**(dB/20)来计算增益因子。
    return audio.apply_gain(gain)


def getBINSize(audio):
    # 创建一个 BytesIO 对象，写入一些数据
    pcm_buffer = io.BytesIO()

    # 获取音频数据
    pcm_data = np.array(audio.get_array_of_samples())

    # 将数据转换为 big-endian 格式
    pcm_data_big_endian = pcm_data.astype(np.int16).byteswap()

    # pcm_data_big_endian数据写入 PCM 文件
    pcm_buffer.write(pcm_data_big_endian.tobytes())

    size_in_bytes = pcm_buffer.getbuffer().nbytes
    return size_in_bytes+160


def outputBIN(audio, outputFile):
    try:
        # 获取音频数据
        pcm_data = np.array(audio.get_array_of_samples())

        # 将数据转换为 big-endian 格式
        pcm_data_big_endian = pcm_data.astype(np.int16).byteswap()

        # 写入 PCM 文件
        with open(outputFile, 'wb') as pcm_file:
            # 写入 160 个 0x00 字节
            pcm_file.write(b'\x00' * 160)
            pcm_file.write(pcm_data_big_endian)

        return 0
    except Exception as e:
        print(f"错误: {e}")
        return -1
